from mmengine.config import read_base
with read_base():
    from .mgam import *

from mgamdata.models.SegFormer3D import SegFormer3D_Encoder_MM, SegFormer3D_Decoder_MM
from mgamdata.models.MoCo import MoCoV3, MoCoV3Head_WithAcc
from mmpretrain.models.losses import CrossEntropyLoss

embed_dims = 32
embed_dims_SegFormer3D = [
    embed_dims, 2*embed_dims, 5*embed_dims, 8*embed_dims]
head_embed_dims = embed_dims_SegFormer3D[-1]//2

model = dict(
    type = MoCoV3,
    backbone_checkpoint=False,
    encoder=dict(
        type=SegFormer3D_Encoder_MM,
        in_channels=in_channels, 
        embed_dims=embed_dims_SegFormer3D,
    ),
    neck=None,
    decoder=dict(
        type=SegFormer3D_Decoder_MM,
        embed_dims=embed_dims_SegFormer3D,
        head_embed_dims=head_embed_dims, 
    ),
    head=dict(
        type=MoCoV3Head_WithAcc, 
        loss=dict(type=CrossEntropyLoss), 
        embed_dim=head_embed_dims, 
        proj_channel=head_embed_dims,
        dim='3d',
    )
)
